import sys

from glob import glob
import os
import re

global_test_func_name_list = []


def modify_content(content, prefix):

    global global_test_func_name_list

    # add tests func for cargo test
    func_pattern = r'\n(\#\[no_mangle\]\n)?(pub\s)?(unsafe\s)?(extern\s\"C\"\s)?fn\s(\w+)\(([^{}]*?)\)([^{]*?)[\{(.*?)\n\}|\{\}]'

    test_func_name_list = []

    matches = re.finditer(func_pattern, content, re.DOTALL)
    for match_item in matches:
        test_func_name_list.append((match_item.group(5), match_item.group(3) is None))
    test_func_name_list = [item for item in test_func_name_list if item[0].startswith(prefix)]

    if len(test_func_name_list) == 0:
        return content
    else:
        global_test_func_name_list += [item[0] for item in test_func_name_list]

    cargo_test_str = (
        "#[cfg(test)]\n"
        "mod tests {\n"
        "    use super::*;\n")
    for test_func_name, is_safe in test_func_name_list:
        if is_safe:
            cargo_test_str += ("    #[test]\n"
                               f"    fn entry_for_{test_func_name}() "
                               "{\n"
                               f"        {test_func_name}();\n"""
                               "    }\n")
        else:
            cargo_test_str += ("    #[test]\n"
                               f"    fn entry_for_{test_func_name}() "
                               "{\n"
                               "        unsafe{\n"
                               f"            {test_func_name}();\n"
                               "        }\n"
                               "    }\n")
    cargo_test_str += "}"
    return content + '\n' + cargo_test_str


def add_cargo_tests(rust_input_dir, rust_output_dir, func_prefix):
    input_file_list = glob(os.path.join(rust_input_dir, '**', '*.*'), recursive=True)
    file_dict = {}
    for filename in input_file_list:
        if os.path.isfile(filename):

            rel_dir_path = os.path.dirname(os.path.relpath(filename, rust_input_dir))
            os.makedirs(os.path.join(rust_output_dir, rel_dir_path), exist_ok=True)

            # if os.path.isfile(os.path.join(rust_input_dir, filename)):
            # print(filename)
            if filename.endswith('.rs') and filename.find('src') != -1:
                with open(filename, 'r') as f:
                    file_content = f.read()
                # print(os.path.relpath(filename, rust_input_dir))
                file_dict[os.path.relpath(filename, rust_input_dir)] = file_content
            else:
                with open(filename, 'rb') as f:
                    file_content = f.read()
                with open(os.path.join(rust_output_dir, os.path.relpath(filename, rust_input_dir)), 'wb') as f:
                    f.write(file_content)

    for filename, file_content in file_dict.items():
        file_content = modify_content(file_content, func_prefix)
        with open(os.path.join(rust_output_dir, filename), 'w') as f:
            f.write(file_content)

    with open(os.path.join(rust_output_dir, 'test_entry_list.txt'), 'w') as f:
        f.write('\n'.join(global_test_func_name_list))


if __name__ == '__main__':
    add_cargo_tests(sys.argv[1], sys.argv[2], sys.argv[3])
